#!/usr/bin/env python3
"""
修复数据库中无效的decimal数据
"""

import os
import sys
import django
from decimal import Decimal, InvalidOperation
import math

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'wealth_backend.settings')
django.setup()

from portfolios.models import Portfolio, PortfolioHolding

def is_valid_decimal(value):
    """检查decimal值是否有效"""
    if value is None:
        return False
    try:
        # 检查是否为NaN或无穷大
        if not value.is_finite():
            return False
        # 尝试量化
        value.quantize(Decimal('0.0001'))
        return True
    except (InvalidOperation, AttributeError, ValueError):
        return False

def fix_portfolio_data():
    """修复Portfolio模型中的无效decimal数据"""
    print("开始修复Portfolio数据...")
    
    portfolios = Portfolio.objects.all()
    fixed_count = 0
    
    for portfolio in portfolios:
        updated = False
        
        # 检查和修复所有decimal字段
        decimal_fields = [
            ('total_amount', Decimal('0.00')),
            ('current_value', Decimal('0.00')),
        ]
        
        for field_name, default_value in decimal_fields:
            field_value = getattr(portfolio, field_name)
            if not is_valid_decimal(field_value):
                setattr(portfolio, field_name, default_value)
                updated = True
                print(f"修复Portfolio {portfolio.id} 的{field_name}")
        
        if updated:
            portfolio.save()
            fixed_count += 1
    
    print(f"Portfolio数据修复完成，共修复 {fixed_count} 条记录")

def fix_portfolio_holding_data():
    """修复PortfolioHolding模型中的无效decimal数据"""
    print("开始修复PortfolioHolding数据...")
    
    holdings = PortfolioHolding.objects.all()
    fixed_count = 0
    
    for holding in holdings:
        updated = False
        
        # 检查和修复所有decimal字段
        decimal_fields = [
            ('shares', Decimal('0.0000')),
            ('cost_price', Decimal('0.0000')),
            ('current_price', Decimal('0.0000')),
            ('cost_amount', Decimal('0.00')),
            ('current_amount', Decimal('0.00')),
        ]
        
        for field_name, default_value in decimal_fields:
            field_value = getattr(holding, field_name)
            if not is_valid_decimal(field_value):
                setattr(holding, field_name, default_value)
                updated = True
                print(f"修复PortfolioHolding {holding.id} 的{field_name}")
        
        if updated:
            holding.save()
            fixed_count += 1
    
    print(f"PortfolioHolding数据修复完成，共修复 {fixed_count} 条记录")

def main():
    print("开始修复数据库中的无效decimal数据...")
    
    try:
        fix_portfolio_data()
        fix_portfolio_holding_data()
        print("所有数据修复完成！")
    except Exception as e:
        print(f"修复过程中出现错误: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()